#!/usr/bin/env python

from functools import reduce
import numpy
import scipy.linalg
from pyscf import scf
from pyscf import gto
from pyscf import mcscf, fci
from pyscf.mcscf import dmet_cas

'''
Triplet and quintet energy gap of Iron-Porphyrin molecule

In this example, we use density matrix embedding theory
(ref. Q Sun, JCTC, 10(2014), 3784) to generate initial guess.
'''

#
# For 3d transition metal, people usually consider the so-called double
# d-shell effects for CASSCF calculation.  Double d-shell here refers to 3d
# and 4d atomic orbitals.  Density matrix embedding theory (DMET) provides a
# method to generate CASSCF initial guess in terms of localized orbitals.
# Given DMET impurity and truncated bath, we can select Fe 3d and 4d orbitals
# and a few entangled bath as the active space.
#


##################################################
#
# Quintet
#
##################################################

mol = gto.Mole()
mol.atom = [
    ['Fe', (0.      , 0.0000  , 0.0000)],
    ['N' , (1.9764  , 0.0000  , 0.0000)],
    ['N' , (0.0000  , 1.9884  , 0.0000)],
    ['N' , (-1.9764 , 0.0000  , 0.0000)],
    ['N' , (0.0000  , -1.9884 , 0.0000)],
    ['C' , (2.8182  , -1.0903 , 0.0000)],
    ['C' , (2.8182  , 1.0903  , 0.0000)],
    ['C' , (1.0918  , 2.8249  , 0.0000)],
    ['C' , (-1.0918 , 2.8249  , 0.0000)],
    ['C' , (-2.8182 , 1.0903  , 0.0000)],
    ['C' , (-2.8182 , -1.0903 , 0.0000)],
    ['C' , (-1.0918 , -2.8249 , 0.0000)],
    ['C' , (1.0918  , -2.8249 , 0.0000)],
    ['C' , (4.1961  , -0.6773 , 0.0000)],
    ['C' , (4.1961  , 0.6773  , 0.0000)],
    ['C' , (0.6825  , 4.1912  , 0.0000)],
    ['C' , (-0.6825 , 4.1912  , 0.0000)],
    ['C' , (-4.1961 , 0.6773  , 0.0000)],
    ['C' , (-4.1961 , -0.6773 , 0.0000)],
    ['C' , (-0.6825 , -4.1912 , 0.0000)],
    ['C' , (0.6825  , -4.1912 , 0.0000)],
    ['H' , (5.0441  , -1.3538 , 0.0000)],
    ['H' , (5.0441  , 1.3538  , 0.0000)],
    ['H' , (1.3558  , 5.0416  , 0.0000)],
    ['H' , (-1.3558 , 5.0416  , 0.0000)],
    ['H' , (-5.0441 , 1.3538  , 0.0000)],
    ['H' , (-5.0441 , -1.3538 , 0.0000)],
    ['H' , (-1.3558 , -5.0416 , 0.0000)],
    ['H' , (1.3558  , -5.0416 , 0.0000)],
    ['C' , (2.4150  , 2.4083  , 0.0000)],
    ['C' , (-2.4150 , 2.4083  , 0.0000)],
    ['C' , (-2.4150 , -2.4083 , 0.0000)],
    ['C' , (2.4150  , -2.4083 , 0.0000)],
    ['H' , (3.1855  , 3.1752  , 0.0000)],
    ['H' , (-3.1855 , 3.1752  , 0.0000)],
    ['H' , (-3.1855 , -3.1752 , 0.0000)],
    ['H' , (3.1855  , -3.1752 , 0.0000)],
]
mol.basis = 'ccpvdz'
mol.verbose = 4
mol.output = 'fepor.out'
mol.spin = 4
mol.symmetry = True
mol.build()

mf = scf.ROHF(mol)
mf = scf.fast_newton(mf)
# ao_labels are keyword(s) to match the AOs generated by mol.ao_labels()
# function.  ao_labels can be a string, or a pattern of regular expression
# or a list of strings/patterns.  The code can detect the format and find
# the AOs to be taken into the active space.
# See also function gto.mole.search_ao_label for the rules of ao_pattern
ao_pattern = ['Fe 3d', 'Fe 4d']
ncas, nelecas, mo = dmet_cas.guess_cas(mf, mf.make_rdm1(), ao_pattern)
mc = mcscf.approx_hessian(mcscf.CASSCF(mf, ncas, nelecas))
mc.kernel(mo)
e_q = mc.e_tot  # -2244.82910509839







##################################################
#
# Triplet
#
##################################################

mol.spin = 2
mol.build(0, 0)  # (0, 0) to avoid dumping input file again

mf = scf.ROHF(mol)
mf = scf.fast_newton(mf)

#
# CAS(8e, 11o)
#
mf = scf.ROHF(mol)
mf = scf.fast_newton(mf)
idx3d4d = mol.search_ao_label(['Fe 3d', 'Fe 4d'])
ncas, nelecas, mo = dmet_cas.guess_cas(mf, mf.make_rdm1(), idx3d)
mc = mcscf.approx_hessian(mcscf.CASSCF(mf, ncas, nelecas))
mc.kernel(mo)
e_t = mc.e_tot  # -2244.81493852189


print('E(T) = %.15g  E(Q) = %.15g  gap = %.15g' % (e_t, e_q, e_t-e_q))

